import cv2
import numpy as np
import glob

# 找棋盘格角点
# 阈值
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
# 棋盘格模板规格
w = 9
h = 6

# 世界坐标系中的棋盘格点,例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
objp = np.zeros((w*h, 3), np.float32)
objp[:, :2] = np.mgrid[0:w, 0:h].T.reshape(-1, 2)

# 储存棋盘格角点的世界坐标和图像坐标对
objpoints = [] # 在世界坐标系中的三维点
imgpoints = [] # 在图像平面的二维点

images = glob.glob('test_pic/*.jpg')
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  #图像二值化

    # 找到棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, (w, h), None)

    # 如果找到足够点对，将其存储起来
    if ret == True:
        cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        objpoints.append(objp)
        imgpoints.append(corners)

        # 将角点在图像上显示
        cv2.drawChessboardCorners(img, (w,  h), corners, ret)
        cv2.imshow('findCorners', img)
        cv2.waitKey(500)
cv2.destroyAllWindows()#删除建立的全部窗口

# 标定
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
# 相机标定 返回标定结果、相机的内参数矩阵、畸变系数、旋转矩阵和平移向量

# 保存相机参数
np.savez('C.npz', mtx=mtx, dist=dist, rvecs=rvecs, tvecs=tvecs)
print("ret:", ret)
print("内参数矩阵:\n", mtx)
print("畸变系数:\n", dist)
print("旋转向量:", rvecs)  # 外参数
print("平移向量:", tvecs)  # 外参数

# 消除畸变
# 读入图像
img = cv2.imread('test_pic/right05.jpg')
# 获取图片的长度
h, w = img.shape[:2]
# 根据尺度因子调节相机矩阵
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))
# 调用remap函数
# 校正畸变图片
mapx, mapy = cv2.initUndistortRectifyMap(mtx, dist, None, newcameramtx, (w, h), 5)
dst = cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR)

# 裁剪图片
x, y, w, h = roi
dst = dst[y:y+h, x:x+w]
cv2.imwrite('calibresult.jpg', dst)

# 反投影误差
total_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2)/len(imgpoints2)
    total_error += error
print("total error: {}".format(total_error/len(objpoints)))



import cv2
import numpy as np
import glob

# 加载相机标定的数据
with np.load('C.npz') as X:
    mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]


def draw(img, corners, imgpts):
    """
    在图片上画出三维坐标轴
    :param img: 图片原数据
    :param corners: 图像平面点坐标点
    :param imgpts: 三维点投影到二维图像平面上的坐标
    :return:
    """
    corner = tuple(corners[0].ravel())
    cv2.line(img, corner, tuple(imgpts[0].ravel()), (255, 0, 0), 5)
    cv2.line(img, corner, tuple(imgpts[1].ravel()), (0, 255, 0), 5)
    cv2.line(img, corner, tuple(imgpts[2].ravel()), (0, 0, 255), 5)
    return img


# 初始化目标坐标系的3D点
objp = np.zeros((6 * 7, 3), np.float32)
objp[:, :2] = np.mgrid[0:7, 0:6].T.reshape(-1, 2)
# 初始化三维坐标系
axis = np.float32([[3, 0, 0], [0, 3, 0], [0, 0, -3]]).reshape(-1, 3)  # 坐标轴
# 加载打包所有图片数据
images = glob.glob('test_pic/*.jpg')
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 找到图像平面点坐标点
    ret, corners = cv2.findChessboardCorners(gray, (7, 6), None)
    if ret:
        # PnP计算得出旋转向量和平移向量
        _, rvecs, tvecs, _ = cv2.solvePnPRansac(objp, corners, mtx, dist)
        print("旋转变量", rvecs)
        print("平移变量", tvecs)
        # 计算三维点投影到二维图像平面上的坐标
        imgpts, jac = cv2.projectPoints(axis, rvecs, tvecs, mtx, dist)
        # 把坐标显示图片上
        img = draw(img, corners, imgpts)
        cv2.imshow('img', img)
        cv2.waitKey(500)

cv2.destroyAllWindows()



